import os  # 用于访问环境变量和操作系统功能

from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import chat_agent_executor
from pydantic import SecretStr

api_key = os.getenv("DASHSCOPE_API_KEY")
# 检查API密钥是否存在，若不存在则抛出错误提示
if not api_key:
    raise ValueError("请设置环境变量DASHSCOPE_API_KEY（阿里云百炼API-KEY）")
# 1. 创建大语言模型实例（适配阿里云百炼的DeepSeek模型）
model = ChatOpenAI(
    model="qwen-plus-latest",  # 模型名称，需与阿里云百炼平台支持的模型名一致
    temperature=0,  # 生成内容的随机性（0表示确定性输出，1表示最大随机性）
    max_tokens=None,  # 最大生成 tokens 数（None表示使用模型默认值）
    timeout=None,  # 超时时间（None表示无超时限制）
    max_retries=2,  # 调用失败时的最大重试次数
    api_key=SecretStr(api_key),  # 以安全字符串形式传递API密钥
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 阿里云百炼的OpenAI兼容接口地址
)

# SQLAIChemy 连接MySQL数据库
HOSTNAME = '127.0.0.1'
PORT = '3306'
DATABASE = 'quickstarts'
USERNAME = 'root'
PASSWORD = '492396'
# mysqlclient 驱动URL
MYSQL_URL = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)
# 获得连接
db = SQLDatabase.from_uri(MYSQL_URL)

# 创建工具类
toolkit = SQLDatabaseToolkit(db=db, llm=model)
tools = toolkit.get_tools()

# 使用Agent完成这个数据库的整合
system_prompt = """
您是一个被设计用来气SQL数据库交互的代理。
给定一个输入问题，创建一个语法正确的SQL语句并执行，然后查看查询结果并返回答案。
除非用户指定了他们想要获得的示例的具体数量，否则始终将SQL查询限制为最多10个结果。
你可以按相关列对结果进行排序，以返回MySQL数据库中最匹配的数据。
您可以使用与数据库交互的工具。在执行查询之前，你必须仔细检査。如果在执行查询时出现错误，请重写查询并重试。
不要对数据库做任何DML语句(插入，更新，删除，删除等)

首先，你应该查看数据库中的表，看看可以查询什么。
不要跳过这一步。
然后查询最相关的表的模式。
"""
system_message = SystemMessage(content=system_prompt)

# 创建代理
agent_executor = chat_agent_executor.create_tool_calling_executor(model, tools, prompt=system_message)

resp = agent_executor.invoke({'messages': [HumanMessage(content='请问哪个民族的人数最多的,并给出SQL语句？')]})

result = resp['messages']
# 一般最后一个是AI回答
print(result[len(result) - 1])